# model/predictor.py

import joblib
import pandas as pd
import os
from config.features import FEATURE_SETS

class ModelPredictor:
    def __init__(self, model_dir="models"):
        self.model_dir = model_dir

    def load_model(self, model_name):
        model_path = os.path.join(self.model_dir, f"lgbm_{model_name}.pkl")
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model file not found: {model_path}")
        return joblib.load(model_path)

    def predict(self, model_name, feature_df):
        model = self.load_model(model_name)
        if feature_df.empty:
            return pd.DataFrame(columns=["score", "prediction"])

        X = feature_df.copy()
        scores = model.predict_proba(X)[:, 1]
        preds = model.predict(X)

        result = feature_df.copy()
        result["score"] = scores
        result["prediction"] = preds
        return result

    def predict_from_dict(self, model_name, feature_dict):
        if model_name not in FEATURE_SETS:
            raise ValueError(f"未知模型名: {model_name}")

        feature_keys = FEATURE_SETS[model_name]

        # 检查所有特征是否齐全
        missing_keys = [key for key in feature_keys if key not in feature_dict]
        if missing_keys:
            print(f"feature_dict: {feature_dict}")
            raise ValueError(f"缺少必要特征: {missing_keys}")

        # 构建 DataFrame
        df = pd.DataFrame([{key: feature_dict[key] for key in feature_keys}])

        # 执行预测
        model = self.load_model(model_name)
        score = model.predict_proba(df)[:, 1][0]
        pred = model.predict(df)[0]

        # ✅ 强制转换 prediction 为 Python 原生 bool 类型
        return {
            "score": float(score),
            "prediction": bool(pred.item())  # 修复 numpy.bool_ 引发的 JSON 序列化问题
        }

